#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
数据库初始化脚本
用于创建数据库和表
"""

# 导入操作系统相关模块
import os
# 导入系统相关模块
import sys
# 导入dotenv用于加载环境变量
from dotenv import load_dotenv
# 导入PyMySQL用于数据库连接
import pymysql
# 导入SQLAlchemy相关组件
from sqlalchemy import create_engine, text
# 导入SQLAlchemy异常处理
from sqlalchemy.exc import SQLAlchemyError

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# 加载环境变量
load_dotenv()

def get_db_config():
    """
    获取数据库配置
    Returns:
        dict: 数据库配置字典
    """
    return {
        'user': os.getenv('MYSQL_USER', 'root'),
        'password': os.getenv('MYSQL_PASS', 'password'),
        'host': os.getenv('MYSQL_HOST', 'localhost'),
        'port': int(os.getenv('MYSQL_PORT', 3306)),
        'db': os.getenv('MYSQL_DB', 'immune_data_platform')
    }

def create_database():
    """
    创建数据库
    如果数据库不存在则创建
    Returns:
        bool: 创建成功返回True，否则返回False
    """
    config = get_db_config()
    print(f"正在连接到MySQL服务器: {config['host']}:{config['port']}")
    
    try:
        # 先连接到MySQL服务器（不指定数据库）
        connection = pymysql.connect(
            host=config['host'],
            port=config['port'],
            user=config['user'],
            password=config['password'],
            charset='utf8mb4'
        )
        
        with connection.cursor() as cursor:
            # 创建数据库
            sql = f"CREATE DATABASE IF NOT EXISTS `{config['db']}` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"
            cursor.execute(sql)
            print(f"数据库 '{config['db']}' 创建成功或已存在")
        
        connection.close()
        return True
        
    except Exception as e:
        print(f"创建数据库时出错: {e}")
        return False

def test_database_connection():
    """
    测试数据库连接
    验证是否可以成功连接到数据库
    Returns:
        bool: 连接成功返回True，否则返回False
    """
    config = get_db_config()
    try:
        # 构建数据库连接字符串
        db_uri = f"mysql+pymysql://{config['user']}:{config['password']}@{config['host']}:{config['port']}/{config['db']}?charset=utf8mb4"
        print(f"尝试连接到数据库: {db_uri}")
        
        # 创建测试引擎
        test_engine = create_engine(db_uri, echo=False)
        
        # 尝试连接
        with test_engine.connect() as connection:
            result = connection.execute(text("SELECT 1"))
            if result.fetchone():
                print("数据库连接成功!")
                return True
    except Exception as e:
        print(f"数据库连接失败: {e}")
        return False

def create_tables():
    """
    创建数据库表
    根据模型定义创建所有数据表
    Returns:
        bool: 创建成功返回True，否则返回False
    """
    print("开始创建数据库表...")
    
    try:
        # 导入模型
        from utils.models import Base, engine, User, File, Workflow, Result, Role, Permission, RolePermission, WorkflowTemplate, WorkflowExecution, ResultShare, SystemConfig, OperationLog
        
        # 创建所有表
        Base.metadata.create_all(bind=engine)
        print("数据库表创建成功!")
        
        # 显示创建的表
        from sqlalchemy import inspect
        inspector = inspect(engine)
        tables = inspector.get_table_names()
        print("已创建的表:")
        for table in sorted(tables):
            print(f"  - {table}")
            
        return True
        
    except SQLAlchemyError as e:
        print(f"SQLAlchemy错误: {e}")
        return False
    except Exception as e:
        print(f"创建数据库表时出错: {e}")
        import traceback
        traceback.print_exc()
        return False

def seed_initial_data():
    """
    初始化基础数据
    创建基础角色和权限数据
    Returns:
        bool: 初始化成功返回True，否则返回False
    """
    print("开始初始化基础数据...")
    
    try:
        from utils.models import SessionLocal, Role, Permission, RolePermission
        from sqlalchemy.orm import sessionmaker
        
        # 创建会话
        db = SessionLocal()
        
        # 创建基础角色
        roles_data = [
            {"name": "admin", "description": "管理员"},
            {"name": "researcher", "description": "研究员"},
        ]
        
        for role_data in roles_data:
            role = db.query(Role).filter(Role.name == role_data["name"]).first()
            if not role:
                role = Role(**role_data)
                db.add(role)
        
        # 创建基础权限
        permissions_data = [
            {"name": "user_create", "description": "创建用户"},
            {"name": "user_read", "description": "读取用户"},
            {"name": "user_update", "description": "更新用户"},
            {"name": "user_delete", "description": "删除用户"},
            {"name": "file_create", "description": "创建文件"},
            {"name": "file_read", "description": "读取文件"},
            {"name": "file_update", "description": "更新文件"},
            {"name": "file_delete", "description": "删除文件"},
            {"name": "workflow_create", "description": "创建工作流"},
            {"name": "workflow_read", "description": "读取工作流"},
            {"name": "workflow_update", "description": "更新工作流"},
            {"name": "workflow_delete", "description": "删除工作流"},
            {"name": "result_create", "description": "创建结果"},
            {"name": "result_read", "description": "读取结果"},
            {"name": "result_update", "description": "更新结果"},
            {"name": "result_delete", "description": "删除结果"},
        ]
        
        for perm_data in permissions_data:
            permission = db.query(Permission).filter(Permission.name == perm_data["name"]).first()
            if not permission:
                permission = Permission(**perm_data)
                db.add(permission)
        
        db.commit()
        
        # 为管理员角色分配所有权限
        admin_role = db.query(Role).filter(Role.name == "admin").first()
        if admin_role:
            permissions = db.query(Permission).all()
            for permission in permissions:
                role_perm = db.query(RolePermission).filter(
                    RolePermission.role_id == admin_role.id,
                    RolePermission.permission_id == permission.id
                ).first()
                if not role_perm:
                    role_perm = RolePermission(
                        role_id=admin_role.id,
                        permission_id=permission.id
                    )
                    db.add(role_perm)
        db.commit()
        db.close()
        print("基础数据初始化完成!")
        return True
        
    except Exception as e:
        print(f"初始化基础数据时出错: {e}")
        import traceback
        traceback.print_exc()
        return False

def main():
    """
    主函数
    执行数据库初始化的完整流程
    """
    print("开始初始化数据库...")
    
    # 1. 创建数据库
    print("步骤1: 创建数据库")
    if not create_database():
        print("创建数据库失败")
        sys.exit(1)
    
    # 2. 测试数据库连接
    print("\n步骤2: 测试数据库连接")
    if not test_database_connection():
        print("数据库连接测试失败")
        sys.exit(1)
    
    # 3. 创建表
    print("\n步骤3: 创建数据表")
    if not create_tables():
        print("创建数据表失败")
        sys.exit(1)
    
    # 4. 初始化基础数据
    print("\n步骤4: 初始化基础数据")
    if not seed_initial_data():
        print("初始化基础数据失败")
        sys.exit(1)
    
    print("\n数据库初始化完成!")

# 程序入口
if __name__ == "__main__":
    main()